"""
测试数据加载器模块
"""
import os
import sys
import unittest
import numpy as np
import torch
from torch.utils.data import DataLoader

# 添加项目根目录到路径
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..')))

from sichuanmajiang.data.data_loader import (
    MahjongDataset, BalancedSampler, DataLoaderBuilder, ParallelDataLoader,
    create_data_loader, create_balanced_data_loader, load_dataset
)
from sichuanmajiang.tests.test_utils import (
    create_test_game_state, create_test_directory, remove_test_directory,
    create_test_features
)


class TestMahjongDataset(unittest.TestCase):
    """
    测试MahjongDataset类
    """
    def setUp(self):
        """
        测试前的设置
        """
        self.test_dir = os.path.join(os.path.dirname(__file__), 'test_data')
        create_test_directory(self.test_dir)
        
        # 创建测试数据
        self.num_samples = 100
        self.feature_dim = 324
        
        # 生成随机特征和标签
        self.features = np.random.rand(self.num_samples, self.feature_dim).astype(np.float32)
        self.actions = np.random.randint(0, 6, self.num_samples)
        self.rewards = np.random.rand(self.num_samples).astype(np.float32)
        
        # 保存测试数据
        self.dataset_path = os.path.join(self.test_dir, 'test_dataset.npz')
        np.savez(
            self.dataset_path,
            features=self.features,
            actions=self.actions,
            rewards=self.rewards
        )
    
    def tearDown(self):
        """
        测试后的清理
        """
        # 移除测试目录
        remove_test_directory(self.test_dir)
    
    def test_init_from_numpy(self):
        """
        测试从NumPy数组初始化
        """
        # 创建数据集
        dataset = MahjongDataset(
            features=self.features,
            actions=self.actions,
            rewards=self.rewards
        )
        
        # 检查数据集大小
        self.assertEqual(len(dataset), self.num_samples)
    
    def test_init_from_file(self):
        """
        测试从文件初始化
        """
        # 创建数据集
        dataset = MahjongDataset.from_file(self.dataset_path)
        
        # 检查数据集大小
        self.assertEqual(len(dataset), self.num_samples)
    
    def test_getitem(self):
        """
        测试获取单个样本
        """
        # 创建数据集
        dataset = MahjongDataset(
            features=self.features,
            actions=self.actions,
            rewards=self.rewards
        )
        
        # 获取样本
        sample_idx = 10
        sample = dataset[sample_idx]
        
        # 检查样本
        self.assertIn('features', sample)
        self.assertIn('action', sample)
        self.assertIn('reward', sample)
        
        # 检查数据类型
        self.assertIsInstance(sample['features'], torch.Tensor)
        self.assertIsInstance(sample['action'], torch.Tensor)
        self.assertIsInstance(sample['reward'], torch.Tensor)
        
        # 检查值是否正确
        np.testing.assert_array_almost_equal(
            sample['features'].numpy(),
            self.features[sample_idx]
        )
        self.assertEqual(sample['action'].item(), self.actions[sample_idx])
        self.assertAlmostEqual(sample['reward'].item(), self.rewards[sample_idx])
    
    def test_len(self):
        """
        测试获取数据集长度
        """
        # 创建数据集
        dataset = MahjongDataset(
            features=self.features,
            actions=self.actions,
            rewards=self.rewards
        )
        
        # 检查长度
        self.assertEqual(len(dataset), self.num_samples)
    
    def test_save_dataset(self):
        """
        测试保存数据集
        """
        # 创建数据集
        dataset = MahjongDataset(
            features=self.features,
            actions=self.actions,
            rewards=self.rewards
        )
        
        # 保存数据集
        save_path = os.path.join(self.test_dir, 'saved_dataset.npz')
        dataset.save_dataset(save_path)
        
        # 检查文件是否存在
        self.assertTrue(os.path.exists(save_path))
        
        # 加载保存的数据集并比较
        loaded_dataset = MahjongDataset.from_file(save_path)
        self.assertEqual(len(loaded_dataset), len(dataset))


class TestBalancedSampler(unittest.TestCase):
    """
    测试BalancedSampler类
    """
    def setUp(self):
        """
        测试前的设置
        """
        # 创建不平衡的标签
        self.num_samples = 100
        # 创建不平衡的动作分布
        self.actions = np.zeros(self.num_samples, dtype=int)
        # 大部分是动作0，少量其他动作
        self.actions[:80] = 0
        self.actions[80:85] = 1
        self.actions[85:90] = 2
        self.actions[90:95] = 3
        self.actions[95:100] = 4
    
    def test_sampler(self):
        """
        测试平衡采样器
        """
        # 创建采样器
        sampler = BalancedSampler(self.actions)
        
        # 获取采样索引
        sample_indices = list(sampler)
        
        # 检查采样数量
        self.assertEqual(len(sample_indices), len(self.actions))
        
        # 检查采样后的分布是否更平衡
        sampled_actions = self.actions[sample_indices]
        action_counts = np.bincount(sampled_actions)
        
        # 各个动作的数量应该大致相等
        max_count = max(action_counts)
        min_count = min(action_counts)
        
        # 最大和最小计数的差异应该不超过10%
        self.assertLessEqual((max_count - min_count) / max_count, 0.3)


class TestDataLoaderBuilder(unittest.TestCase):
    """
    测试DataLoaderBuilder类
    """
    def setUp(self):
        """
        测试前的设置
        """
        self.test_dir = os.path.join(os.path.dirname(__file__), 'test_data')
        create_test_directory(self.test_dir)
        
        # 创建测试数据
        self.num_samples = 100
        self.feature_dim = 324
        
        # 生成随机特征和标签
        self.features = np.random.rand(self.num_samples, self.feature_dim).astype(np.float32)
        self.actions = np.random.randint(0, 6, self.num_samples)
        self.rewards = np.random.rand(self.num_samples).astype(np.float32)
        
        # 创建数据集
        self.dataset = MahjongDataset(
            features=self.features,
            actions=self.actions,
            rewards=self.rewards
        )
    
    def tearDown(self):
        """
        测试后的清理
        """
        # 移除测试目录
        remove_test_directory(self.test_dir)
    
    def test_build_default(self):
        """
        测试默认构建
        """
        # 创建构建器
        builder = DataLoaderBuilder(self.dataset)
        
        # 构建数据加载器
        data_loader = builder.build()
        
        # 检查类型
        self.assertIsInstance(data_loader, DataLoader)
        
        # 检查批次大小
        self.assertEqual(data_loader.batch_size, 32)  # 默认批次大小
    
    def test_build_with_custom_params(self):
        """
        测试自定义参数构建
        """
        # 创建构建器并设置参数
        builder = DataLoaderBuilder(self.dataset)
        builder.set_batch_size(64)
        builder.set_shuffle(True)
        builder.set_num_workers(2)
        
        # 构建数据加载器
        data_loader = builder.build()
        
        # 检查参数是否正确设置
        self.assertEqual(data_loader.batch_size, 64)
        self.assertTrue(data_loader.shuffle)
        self.assertEqual(data_loader.num_workers, 2)
    
    def test_build_balanced(self):
        """
        测试构建平衡的数据加载器
        """
        # 创建构建器
        builder = DataLoaderBuilder(self.dataset)
        
        # 构建平衡的数据加载器
        data_loader = builder.build_balanced()
        
        # 检查类型
        self.assertIsInstance(data_loader, DataLoader)
        
        # 检查是否使用了采样器
        self.assertIsNotNone(data_loader.sampler)
    
    def test_build_multiple(self):
        """
        测试构建多个数据加载器
        """
        # 创建构建器
        builder = DataLoaderBuilder(self.dataset)
        
        # 构建训练和验证数据加载器
        train_loader, val_loader = builder.build_train_val_split(val_ratio=0.2)
        
        # 检查类型
        self.assertIsInstance(train_loader, DataLoader)
        self.assertIsInstance(val_loader, DataLoader)
        
        # 检查数据量比例
        train_size = sum(1 for _ in train_loader) * train_loader.batch_size
        val_size = sum(1 for _ in val_loader) * val_loader.batch_size
        total_size = train_size + val_size
        
        # 应该接近总数
        self.assertLessEqual(abs(total_size - self.num_samples), train_loader.batch_size)


class TestParallelDataLoader(unittest.TestCase):
    """
    测试ParallelDataLoader类
    """
    def setUp(self):
        """
        测试前的设置
        """
        self.test_dir = os.path.join(os.path.dirname(__file__), 'test_data')
        create_test_directory(self.test_dir)
        
        # 创建测试数据
        self.num_samples = 100
        self.feature_dim = 324
        
        # 生成随机特征和标签
        self.features = np.random.rand(self.num_samples, self.feature_dim).astype(np.float32)
        self.actions = np.random.randint(0, 6, self.num_samples)
        self.rewards = np.random.rand(self.num_samples).astype(np.float32)
        
        # 创建数据集
        self.dataset = MahjongDataset(
            features=self.features,
            actions=self.actions,
            rewards=self.rewards
        )
    
    def tearDown(self):
        """
        测试后的清理
        """
        # 移除测试目录
        remove_test_directory(self.test_dir)
    
    def test_parallel_loader(self):
        """
        测试并行数据加载器
        """
        # 创建并行数据加载器
        parallel_loader = ParallelDataLoader(
            self.dataset,
            batch_size=32,
            num_workers=2,
            num_chunks=2
        )
        
        # 获取批次数量
        batch_count = 0
        for batch in parallel_loader:
            batch_count += 1
            
            # 检查批次内容
            self.assertIn('features', batch)
            self.assertIn('action', batch)
            self.assertIn('reward', batch)
        
        # 应该有多个批次
        self.assertGreater(batch_count, 0)
    
    def test_context_manager(self):
        """
        测试上下文管理器功能
        """
        # 使用上下文管理器
        with ParallelDataLoader(
            self.dataset,
            batch_size=32,
            num_workers=2
        ) as loader:
            # 获取一个批次
            batch = next(iter(loader))
            
            # 检查批次
            self.assertIn('features', batch)


class TestHelperFunctions(unittest.TestCase):
    """
    测试辅助函数
    """
    def setUp(self):
        """
        测试前的设置
        """
        self.test_dir = os.path.join(os.path.dirname(__file__), 'test_data')
        create_test_directory(self.test_dir)
        
        # 创建测试数据
        self.num_samples = 100
        self.feature_dim = 324
        
        # 生成随机特征和标签
        self.features = np.random.rand(self.num_samples, self.feature_dim).astype(np.float32)
        self.actions = np.random.randint(0, 6, self.num_samples)
        self.rewards = np.random.rand(self.num_samples).astype(np.float32)
        
        # 保存测试数据
        self.dataset_path = os.path.join(self.test_dir, 'test_dataset.npz')
        np.savez(
            self.dataset_path,
            features=self.features,
            actions=self.actions,
            rewards=self.rewards
        )
    
    def tearDown(self):
        """
        测试后的清理
        """
        # 移除测试目录
        remove_test_directory(self.test_dir)
    
    def test_create_data_loader(self):
        """
        测试create_data_loader函数
        """
        # 创建数据加载器
        loader = create_data_loader(
            features=self.features,
            actions=self.actions,
            rewards=self.rewards,
            batch_size=32
        )
        
        # 检查类型
        self.assertIsInstance(loader, DataLoader)
    
    def test_create_balanced_data_loader(self):
        """
        测试create_balanced_data_loader函数
        """
        # 创建平衡数据加载器
        loader = create_balanced_data_loader(
            features=self.features,
            actions=self.actions,
            rewards=self.rewards,
            batch_size=32
        )
        
        # 检查类型
        self.assertIsInstance(loader, DataLoader)
        
        # 检查是否使用了采样器
        self.assertIsNotNone(loader.sampler)
    
    def test_load_dataset(self):
        """
        测试load_dataset函数
        """
        # 加载数据集
        dataset = load_dataset(self.dataset_path)
        
        # 检查类型
        self.assertIsInstance(dataset, MahjongDataset)
        
        # 检查长度
        self.assertEqual(len(dataset), self.num_samples)


if __name__ == '__main__':
    unittest.main()
